#include "stdafx.h"
#include "CertificateGenerator.h"

#include "TBSCertificateGenerator.h"
#include "ASN1Null.h"

#include <osrng.h>
#include <rsa.h>
#include <hex.h>     // Hex Encoder/Decoder
#include <md5.h>

#include "Utils.h"

using namespace CryptoPP;

namespace AppliedCryptography
{
	/*
	CertificateGenerator::CertificateGenerator)
		: _tbsGen(new TBSCertificateGenerator()),
		_sigAlgOid(new ASN1ObjectIdentifier()),
		_signature(new ASN1BitString())
	{
	};
	*/

	CertificateGenerator::CertificateGenerator(
		TBSCertificateGenerator* tbsGen, 
		ASN1ObjectIdentifier* sigAlgOID,
		ASN1BitString* signature)
		: _tbsGen(tbsGen),
		_sigAlgOid(sigAlgOID),
		_signature(signature)
	{
		/* nasing */
	};

	CertificateGenerator::~CertificateGenerator()
	{
		delete _tbsGen;
		delete _sigAlgOid;
		delete _signature;
	};

	ASN1Sequence& CertificateGenerator::add(ASN1Encodable& obj)
	{
		this->ASN1Sequence::add(obj);
		return *this;
	};

	ASN1Encodable& CertificateGenerator::get(unsigned int index) const
	{
		return this->ASN1Sequence::get(index);
	};

	unsigned int CertificateGenerator::size() const
	{
		return -1;
	};

	CertificateGenerator& CertificateGenerator::setSerialNumber(unsigned long serialNumber)
	{
		_tbsGen->setSerialNumber(serialNumber);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setIssuer(Name& issuer)
	{
		_tbsGen->setIssuer(issuer);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setNotBefore(ASN1UTCTime& startTime)
	{
		_tbsGen->setStartDate(startTime);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setNotAfter(ASN1UTCTime& endTime)
	{
		_tbsGen->setEndDate(endTime);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setSubject(Name& subject)
	{
		_tbsGen->setSubject(subject);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setPublicKey(ASN1Encoded& key)
	{
		_tbsGen->setPublicKey(key);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setExtensions(Extensions& extensions)
	{
		_tbsGen->setExtensions(extensions);
		return *this;
	};

	CertificateGenerator& CertificateGenerator::setSignatureAlgorithm(ASN1ObjectIdentifier& algOID)
	{
		*_sigAlgOid = algOID;
		_tbsGen->setSignatureAlgorithm(algOID);
		return *this;
	};

	std::vector<byte> CertificateGenerator::getContentBits() const
	{
		//throw std::exception("Not implemented!");
		ASN1Sequence seq;
		ASN1Sequence algSig;
		ASN1Null null;

		algSig.add(*_sigAlgOid);
		algSig.add(null);

		seq.add(*_tbsGen);
		seq.add(algSig);
		seq.add(*_signature);

		return seq.getContentBits();
	};


	CertificateGenerator& CertificateGenerator::generateSignature(ASN1Encoded& privateKey)
	{
		std::vector<byte> tbs = _tbsGen->encode();
		byte* keyBytes = new byte[privateKey.size()];
		std::string keyString;
		RSA::PrivateKey rsaKey;
		AutoSeededRandomPool rng;
		byte* data = new byte[tbs.size()];

		for(unsigned int i=0; i<tbs.size(); ++i)
		{
			data[i] = tbs[i];
		};

		for(unsigned int i=0; i<privateKey.size(); ++i)
		{
			keyBytes[i] = privateKey.encode()[i];
		};

		HexEncoder hexEnc(new StringSink(keyString));
		hexEnc.Put(keyBytes, privateKey.size());

		HexDecoder hexDec;
		hexDec.Put((byte*)keyString.c_str(), keyString.size());
		hexDec.MessageEnd();

		rsaKey.Load(hexDec);

		//Weak::RSASSA_PKCS1v15_MD5_Signer signer(rsaKey);
		RSASSA_PKCS1v15_SHA_Signer signer(rsaKey);
		// Create signature space
		size_t length = signer.MaxSignatureLength();
		byte* signature = new byte[length];

		// Sign message
		unsigned int sigSize = signer.SignMessage( rng, data,
			tbs.size(), signature );

		_signature->set(signature, sigSize);
		std::string sigHex = Vector2HexStr(_signature->get());

		delete [] data;
		delete [] keyBytes;
		delete [] signature;

		//throw std::exception("Not implemented!");
		return *this;
	};
}; // namespace AppliedCryptography
